/*
 * Kalman filter for Flying Fox
 *
 * from paparazzi autopilot
 * also by Zik Saleeba 2008-04-05
 */

#include <string.h>
#include <math.h>

#include "FreeRTOS.h"
#include "kalman.h"
#include "matrix.h"


/*
 * ukf_filter_new
 */

ukf_filter
ukf_filter_new(unsigned int state_dim, 
               unsigned int measure_dim, 
               double *Q,
               double *R,
               filter_function ffun,
               measure_function mfun) {
    ukf_filter filter;
    int Size;
    unsigned err = 0;
    // nothing to do if no state or measurement !
    if(state_dim == 0 || measure_dim == 0)
        return 0;
    // alloc new structure
    filter = pvPortMalloc(sizeof(struct ukf_filter_t));
    // returns 0 if allocation fails
    if(filter == 0) 
        return 0;
    // fills the structure
    filter->state_dim = state_dim;
    filter->measure_dim = measure_dim;
    filter->ffun = ffun;
    filter->mfun = mfun;
    
    filter->x = pvPortMalloc(state_dim * sizeof(double));
    err |= (filter->x == 0);
    
    filter->y = pvPortMalloc(measure_dim * sizeof(double));
    err |= (filter->y == 0);
    
    Size = state_dim * state_dim;
    
    filter->P = pvPortMalloc(Size * sizeof(double));
    err |= (filter->P == 0);
    
    Size = 2 * state_dim + 1;
    
    filter->wm = pvPortMalloc(Size * sizeof(double));
    err |= (filter->wm == 0);
    
    filter->wc = pvPortMalloc(Size * sizeof(double));
    err |= (filter->wc == 0);
    
    filter->sigma_point = pvPortMalloc(Size * state_dim * sizeof(double));
    err |= (filter->sigma_point == 0);
    
    Size = filter->state_dim;
    
    filter->sigma = pvPortMalloc(Size * sizeof(double));
    err |= (filter->sigma == 0);
    
    filter->PM = pvPortMalloc(Size * Size * sizeof(double));
    err |= (filter->PM == 0);
    
    filter->PM_save = pvPortMalloc(Size * Size * sizeof(double));
    err |= (filter->PM == 0);
    
    filter->xm = pvPortMalloc(Size * sizeof(double));
    err |= (filter->xm == 0);
    
    filter->ym = pvPortMalloc(filter->measure_dim * sizeof(double));
    err |= (filter->ym == 0);
    
    Size = 2 * filter->state_dim + 1;
    filter->khi = pvPortMalloc(Size * filter->state_dim * sizeof(double));
    err |= (filter->khi == 0);
    
    filter->khi_y = pvPortMalloc(Size * filter->measure_dim * sizeof(double));
    err |= (filter->khi_y == 0);
    
    Size = filter->measure_dim;
    filter->Pyy = pvPortMalloc(Size * Size * sizeof(double));
    err |= (filter->Pyy == 0);
    
    filter->Pxy = pvPortMalloc(Size * filter->state_dim * sizeof(double));
    err |= (filter->Pxy == 0);
     
    
    filter->dx = pvPortMalloc(filter->state_dim * sizeof(double));
    err |= (filter->dx == 0);
    
    filter->dy = pvPortMalloc(filter->measure_dim * sizeof(double));
    err |= (filter->dy == 0);
    
    filter->gain = pvPortMalloc(filter->state_dim * filter->measure_dim * sizeof(double));
    err |= (filter->gain == 0);
    
    filter->sigma_y = pvPortMalloc(filter->measure_dim * sizeof(double));
    err |= (filter->sigma_y == 0);
    
    filter->KL = pvPortMalloc(filter->state_dim * filter->measure_dim * sizeof(double));
    err |= (filter->KL == 0);
        
    if(err != 0) {
        vPortFree(filter);
        return 0;
    }
    
    Size = filter->state_dim;
    filter->Q = pvPortMalloc(Size * Size * sizeof(double));
    if(filter->Q == 0) {
        vPortFree(filter);
        return 0;
    }
    memcpy(filter->Q, Q, Size * Size * sizeof(double));
    
    Size = filter->measure_dim;
    filter->R = pvPortMalloc(Size * Size * sizeof(double));
    if(filter->R == 0) {
        vPortFree(filter);
        return 0;
    }
    memcpy(filter->R, R, Size * Size * sizeof(double));
    
    // returns the newly allocated structure
    return filter;
}

/*
 * ukf_filter_compute_weights
 */

void
ukf_filter_compute_weights(ukf_filter filter,
                     double alpha, 
                     double ZCount, 
                     double beta) {
    double l;
    double lam;
    unsigned YCount;
    
    if(filter == 0) return;
    l = (double)filter->state_dim;
    // lambda parameter
    lam = alpha * alpha *( l + ZCount) - l;
    
    filter->wm[0] = lam / (lam + l);
    filter->wc[0] = filter->wm[0] + (1.0 - alpha * alpha + beta);
    for(YCount = 1 ; YCount <= 2*filter->state_dim ; YCount++) {
        filter->wm[YCount] = 0.5 / (lam + l);
        filter->wc[YCount] = 0.5 / (lam + l);
    }
    filter->gamma = alpha*sqrt(l + ZCount);
}


/*
 * ukf_filter_reset
 */

void 
ukf_filter_reset(ukf_filter filter,
                 double *x0,
                 double *P0) {
    if(filter != 0) {
        // state of the filter 
        memcpy(filter->x, x0, filter->state_dim * sizeof(double));
        memcpy(filter->P, P0,
               filter->state_dim * filter->state_dim * sizeof(double));
    }
}

/*
 * ukf_filter_get_state
 */

void
ukf_filter_get_state(ukf_filter filter, double *x, double* P){
    if(filter != 0) {
        memcpy(x, filter->x, filter->state_dim * sizeof(double));
        memcpy(P, filter->P, filter->state_dim * filter->state_dim * sizeof(double));
    }
}


/*
 * ukf_filter_update
 */

void
ukf_filter_update(ukf_filter filter, double *y, double *u) {
    int l = filter->state_dim;
    int m = filter->measure_dim;
    int i,j,k;
    double t;
    
    // cholesky decomposition of the state covariance matrix
    
    ukf_cholesky_decomposition(filter->P, l, filter->sigma);
    
    //=================================
    // compute sigma points
    for(j = 0 ; j < l ; j++)
        filter->sigma_point[j]= filter->x[j];
    for(i = 0 ; i < l ; i++) {
        for(j = 0 ; j < i ; j++) {
            filter->sigma_point[(i + 1) * l + j] = filter->x[j] ;
            filter->sigma_point[(i + 1 + l) * l + j] = filter->x[j] ;
        }
        filter->sigma_point[(i + 1) * l + i] = filter->x[i] + filter->gamma * filter->sigma[i];
        filter->sigma_point[(i + 1 + l) * l + i] = filter->x[i] - filter->gamma * filter->sigma[i];
        for(j = i + 1 ; j < l ; j++) {
            filter->sigma_point[(i + 1) * l + j] = filter->x[j]  + filter->gamma * filter->P[j * l + i];
            filter->sigma_point[(i + 1 + l) * l + j] = filter->x[j] - filter->gamma * filter->P[j * l + i];
        }
    }
    
    
    //=================================
    // propagate sigma points
    for(i = 0 ; i < 2 * l + 1 ; i++) {
        filter->ffun(&(filter->khi[i * l]) , &(filter->sigma_point[i * l]) , u);
    }
    
    // compute state prediction xm
    for(i = 0 ; i < l ; i++) {
        filter->xm[i] = filter->wm[0] * filter->khi[i];
        for(j = 1 ; j < 2 * l + 1 ; j++) {
            filter->xm[i] += filter->wm[j] * filter->khi[j * l + i];
        }
    }
        
    //================================
    // time update
    
    // start with state covariance matrix
    for(i = 0 ; i < l * l ; i++)
        filter->PM[i] = filter->Q[i];
        
    // accumulate covariances
    for(i = 0 ; i < 2 * l + 1 ; i++) {
        for(j = 0 ; j < l ; j++)
            filter->dx[j]= filter->khi[i * l + j] - filter->xm[j];
        for(j = 0 ; j < l ; j++) {
            for(k = 0 ; k < l ; k++) {
                filter->PM[j * l + k] += filter->wc[i] * filter->dx[j] * filter->dx[k];
            }
        } 
    }
    
    // save PM matrix
    
    for(i = 0 ; i < l * l ; i++)
        filter->PM_save[i] = filter->PM[i];
    
    // redraw sigma points
    
    ukf_cholesky_decomposition(filter->PM, l, filter->sigma);
   
    for(j = 0 ; j < l ; j++)
        filter->sigma_point[j]= filter->xm[j];
    for(i = 0 ; i < l ; i++) {
        for(j = 0 ; j < i ; j++) {
            filter->sigma_point[(i + 1) * l + j] = filter->xm[j] ;
            filter->sigma_point[(i + 1 + l) * l + j] = filter->xm[j] ;
        }
        filter->sigma_point[(i + 1) * l + i] = filter->xm[i] + filter->gamma * filter->sigma[i];
        filter->sigma_point[(i + 1 + l) * l + i] = filter->xm[i] - filter->gamma * filter->sigma[i];
        for(j = i + 1 ; j < l ; j++) {
            filter->sigma_point[(i + 1) * l + j] = filter->xm[j] + filter->gamma * filter->PM[j * l + i];
            filter->sigma_point[(i + 1 + l) * l + j] = filter->xm[j] - filter->gamma * filter->PM[j * l + i];
        }
    }

    

    //=================================
    // propagate measurement
    
    for(i = 0 ; i < 2 * l + 1 ; i++) {
        filter->mfun(&(filter->khi_y[i * m]) , &(filter->sigma_point[i * l]) );
    }

    // measurement prediction   
    
    for(i = 0 ; i < m ; i++) {
        filter->ym[i] = filter->wm[0] * filter->khi_y[i];
        for(j = 1 ; j < 2 * l + 1 ; j++)
            filter->ym[i] += filter->wm[j] * filter->khi_y[j * m + i];
    }
    
    // measurement update
    
    // Pyy matrix
    // start with measure covariance matrix
    
    for(i = 0 ; i < m * m ; i++)
        filter->Pyy[i] = filter->R[i];
    
    // accumulate covariances
    
    for(i = 0 ; i < 2 * l + 1 ; i++) {
        for(j = 0 ; j < m ; j++)
            filter->dy[j]= filter->khi_y[i * m + j] - filter->ym[j];
        for(j = 0 ; j < m ; j++)
            for(k = 0 ; k < m ; k++) {
                filter->Pyy[j * m + k] += filter->wc[i] * filter->dy[j] * filter->dy[k];
            }
    }
    
     // Pxy matrix
    
    for(i = 0 ; i < m * l ; i++)
        filter->Pxy[i] = 0.0;
    
    // accumulate covariances
    
    for(i = 0 ; i < 2 * l + 1 ; i++) {
        for(j = 0 ; j < m ; j++) {
            filter->dy[j]= filter->khi_y[i * m + j] - filter->ym[j];
        }
        for(j = 0 ; j < l ; j++) {
            filter->dx[j] = filter->sigma_point[i * l + j] - filter->xm[j];
        }
        for(j = 0 ; j < l ; j++)
            for(k = 0 ; k < m ; k++) {
                filter->Pxy[j * m + k] += filter->wc[i] * filter->dx[j] * filter->dy[k];
            }
    }
   
    // gain de kalman
    
    ukf_cholesky_decomposition(filter->Pyy, m, filter->sigma_y);
    ukf_cholesky_solve(filter->Pyy, m, filter->sigma_y, filter->Pxy, l, filter->gain);
    
    // restore PM matrix
    
    for(i = 0 ; i < l * l ; i++)
        filter->P[i]= filter->PM_save[i];
    
    // update state
    
    for(j = 0 ; j < m ; j++)
        filter->dy[j] = y[j] - filter->ym[j];
    for(i = 0 ; i < l ; i++) {
        filter->x[i] = filter->xm[i];
        t = 0.0;
        for(j = 0 ; j < m ; j++)
            t += filter->gain[i * m + j] * filter->dy[j];
        filter->x[i] += t;
    }
    
    for(i = 0 ; i < l ; i++)
        for(j = 0 ; j < m ; j++) {
            t = 0.0;
            for(k = j ; k < m ; k++)
                t += filter->gain[i * m + k] * filter->Pyy[k * m + j];
            filter->KL[i * m + j] = t;
        }
    for(i = 0 ; i < l ; i++)
        for(j = 0 ; j < l ; j++) {
            t = 0.0;
            for(k = 0 ; k < m ; k++)
                t += filter->KL[i * m + k] * filter->KL[j * m + k];
            filter->P[i * l + j ] -= t;
        }
   
   // finished with kalman iteration !
    
}
 
